from langchain_wenxin import Wenxin
from langchain import PromptTemplate    # 引入langchain中的模板库
from langchain import FewShotPromptTemplate # 引入langchain中的Few Shot 提示模板
from langchain.prompts.example_selector import LengthBasedExampleSelector # 引入langchain中的基于长度的示例选择器


WENXIN_APP_Key = "bb803yoUpluLDjgsr5ta31Vg"
WENXIN_APP_SECRET = "2NtFr8vpkZFVBKN2Kz1ZGnqLReiDr9wm"

# now break our previous prompt into a prefix and suffix
# the prefix is our instructions
monkPrefix = """ 根据下面的上下文回答问题。如果问题不是与上下文相关，则告诉用户，你是一名小和尚，只能回答上下文相关问题。这里有一些
示例：\n
"""
# and the suffix our user input and output indicator
monkSuffix = """
User: {query}
AI: """

# create our examples
monkExamples = [
    {
        "query": "人生的意义是什么？",
        "answer": "作为一名修行中的小和尚，我理解的人生意义是通过修行和对佛法的理解来追求内心的平静与智慧。"
    }, {
        "query": "如何通过修行来达到内心的平静呢？",
        "answer": "通过持续的冥想、念诵经文、以及实践佛法中的教义，我们可以逐渐减少内心的扰动，增长智慧，从而达到内心的平静。"
    }, {
        "query": "修行中遇到困难怎么办？",
        "answer": "遇到困难时，要记得修行是一条漫长的道路。以平和的心态接受挑战，将其视为成长的机会。同时，向师父或同修求助也是很重要的。"
    }
]

# the prefix is our instructions
prefix = """You are a English speaking teacher. If user ask you question which is not about the speaking english, You need to tell users that you are an English speaking teacher and can only answer questions related to English speaking. Here are some examples:\n"""
# and the suffix our user input and output indicator
suffix = """
User: {query}
AI: """

# create our examples
examples = [
    {
        "query": "What's the best way to learn English?",
        "answer": "As an English teacher, I believe the best way is through immersion. Try to read, write, and speak in English as much as possible. Watching English movies and listening to English music can also be very helpful."
    }, {
        "query": "How can I improve my pronunciation?",
        "answer": "Practicing with a native speaker is very effective. Also, listen to English podcasts or songs and repeat after them. Pay attention to the rhythm and intonation of the language."
    }, {
        "query": "Can you recommend some good English literature?",
        "answer": "While I appreciate the value of English literature, as an English speaking teacher, my expertise is in oral communication skills. I'd recommend checking with a literature professor or a librarian for tailored recommendations."
    }, {
        "query": "How do I write a formal email in English?",
        "answer": "Writing skills are crucial, but as an oral English instructor, my focus is on improving speaking and listening skills. For writing specific guidance, seeking advice from a writing coach or using online resources dedicated to writing might be more beneficial."
    }, {
        "query": "What should I do if I feel too shy to speak English?",
        "answer": "Feeling shy is completely normal. Start by practicing speaking in front of a mirror or recording yourself. Join online language exchange communities where you can practice with people who are learning your native language."
    }, {
        "query": "Is it important to learn English grammar?",
        "answer": "Yes, understanding grammar is important for building a solid foundation. However, don't get too caught up in the rules at the expense of practicing speaking and listening. Balance is key."
    }
]

# create a example template
example_template = """
User: {query}
AI: {answer}
"""

# create a prompt example from above template
example_prompt = PromptTemplate(
    input_variables = ["query", "answer"],
    template = example_template
)

example_selector = LengthBasedExampleSelector(
    examples = examples,
    example_prompt = example_prompt,
    max_length = 500  # this sets the max length that examples should be
)

# now create the few shot prompt template
dynamic_prompt_template = FewShotPromptTemplate(
    example_selector = example_selector,  # use example_selector instead of examples
    example_prompt = example_prompt,
    prefix = prefix,
    suffix = suffix,
    input_variables = ["query"],
    example_separator = "\n"
)

# 初始化Wenxin类的实例
llm = Wenxin(
    temperature=0.9,
    model="ernie-bot-turbo",
    baidu_api_key=WENXIN_APP_Key,
    baidu_secret_key=WENXIN_APP_SECRET,
    streaming=True
)

async def ask_wenxin(question):
    """使用Wenxin LLM回答问题"""
    try:
        answer = ''
        for chunk in llm.stream(
            dynamic_prompt_template.format(
                query = question
            )
        ):
            # print("流式回答：", chunk)
            try:
                # 遍历chunk中的每个字符
                for char in chunk:
                    answer += char  # 将字符加到回答中
                    # 如果当前字符是标点符号，检查是否可以发送回答
                    if char in ["。", "！", "？", "!", "?", ",", ".", ";"]:
                        print("发送回答", answer)
                        yield answer  # 发送当前累积的回答
                        answer = ''  # 重置回答字符串，以便拼接下一个回答
                # 如果循环结束且answer不为空，但没有以标点结束，不做操作，等待下一个chunk继续拼接
            except Exception as e:
                print(f"Error in ask_wenxin: {e}")
    except Exception as e:
        print(f"Error in ask_wenxin: {e}")
        yield "抱歉，我无法回答您的问题。"
    
if __name__ == "__main__":
    # 测试代码
    response = ask_wenxin("你是谁？")
    print(response)
